"""
文件名: database.py
描述: 数据库连接和会话管理模块

本模块包含以下主要功能：
1. PostgreSQL异步数据库连接管理
2. SQLAlchemy异步会话工厂
3. 数据库初始化和迁移支持
4. 连接池配置和管理
5. 事务管理和错误处理

依赖模块:
   - sqlalchemy: ORM框架
   - asyncpg: PostgreSQL异步驱动
   - contextlib: 上下文管理器
   - typing: 类型注解

使用示例:
   >>> from src.core.database import get_db, DatabaseManager
   >>> async with get_db() as db:
   ...     result = await db.execute("SELECT 1")
   >>> 
   >>> # 或者在FastAPI依赖注入中使用
   >>> @app.get("/users")
   >>> async def get_users(db: AsyncSession = Depends(get_db)):
   ...     return await db.execute(select(User))

注意事项:
   - 使用异步数据库操作，避免阻塞
   - 正确管理数据库会话的生命周期
   - 生产环境需要配置连接池参数
   - 定期检查数据库连接健康状态

作者: AI助手
创建日期: 2024-06-27
最后修改: 2024-06-27
版本: 1.0.0
"""

from contextlib import asynccontextmanager
from typing import AsyncGenerator, Optional

from sqlalchemy import create_engine, event, pool
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
    AsyncEngine,
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool

from .config import get_settings
from .logger import get_logger

logger = get_logger(__name__)
settings = get_settings()

# 创建基础模型类
Base = declarative_base()

# 全局数据库引擎和会话工厂
async_engine: Optional[AsyncEngine] = None
async_session_factory: Optional[async_sessionmaker[AsyncSession]] = None


class DatabaseManager:
    """数据库管理器类"""
    
    def __init__(self):
        self.engine: Optional[AsyncEngine] = None
        self.session_factory: Optional[async_sessionmaker[AsyncSession]] = None
        self._initialized = False
    
    async def initialize(self) -> None:
        """
        初始化数据库连接
        
        配置包括：
        1. 创建异步数据库引擎
        2. 配置连接池参数
        3. 设置会话工厂
        4. 注册事件监听器
        """
        if self._initialized:
            logger.warning("数据库已经初始化，跳过重复初始化")
            return
        
        try:
            # 创建异步数据库引擎
            self.engine = create_async_engine(
                settings.database_url,
                # 连接池配置
                poolclass=QueuePool,
                pool_size=10,  # 连接池大小
                max_overflow=20,  # 最大溢出连接数
                pool_pre_ping=True,  # 连接前测试连接是否有效
                pool_recycle=3600,  # 连接回收时间（秒）
                # 查询配置
                echo=settings.debug,  # 是否打印SQL语句
                echo_pool=settings.debug,  # 是否打印连接池信息
                # 其他配置
                future=True,
                connect_args={
                    "server_settings": {
                        "jit": "off",  # 关闭JIT编译，提高连接速度
                    }
                },
            )
            
            # 创建会话工厂
            self.session_factory = async_sessionmaker(
                bind=self.engine,
                class_=AsyncSession,
                expire_on_commit=False,  # 提交后不过期对象
                autoflush=True,  # 自动刷新
                autocommit=False,  # 手动提交
            )
            
            # 设置全局变量
            global async_engine, async_session_factory
            async_engine = self.engine
            async_session_factory = self.session_factory
            
            # 注册数据库事件监听器
            @event.listens_for(self.engine.sync_engine, "connect", once=True)
            def set_database_params(dbapi_connection, connection_record):
                """设置数据库连接参数"""
                if "postgresql" in settings.database_url:
                    # PostgreSQL 连接配置
                    with dbapi_connection.cursor() as cursor:
                        # 设置时区
                        cursor.execute("SET timezone TO 'UTC'")
                        # 设置字符编码
                        cursor.execute("SET client_encoding TO 'UTF8'")
            
            self._initialized = True
            logger.info("数据库初始化成功", database_url=settings.database_url.split("@")[-1])
            
        except Exception as e:
            logger.error(f"数据库初始化失败: {e}", error=str(e))
            raise
    
    async def close(self) -> None:
        """关闭数据库连接"""
        if self.engine:
            await self.engine.dispose()
            logger.info("数据库连接已关闭")
    
    async def create_tables(self) -> None:
        """创建数据库表"""
        if not self.engine:
            raise RuntimeError("数据库引擎未初始化")
        
        try:
            async with self.engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
            logger.info("数据库表创建成功")
        except Exception as e:
            logger.error(f"数据库表创建失败: {e}", error=str(e))
            raise
    
    async def drop_tables(self) -> None:
        """删除数据库表（谨慎使用）"""
        if not self.engine:
            raise RuntimeError("数据库引擎未初始化")
        
        try:
            async with self.engine.begin() as conn:
                await conn.run_sync(Base.metadata.drop_all)
            logger.info("数据库表删除成功")
        except Exception as e:
            logger.error(f"数据库表删除失败: {e}", error=str(e))
            raise
    
    async def health_check(self) -> bool:
        """
        数据库健康检查
        
        Returns:
            bool: 数据库是否健康
        """
        if not self.engine:
            return False
        
        try:
            async with self.engine.connect() as conn:
                await conn.execute("SELECT 1")
            return True
        except Exception as e:
            logger.error(f"数据库健康检查失败: {e}", error=str(e))
            return False


# 全局数据库管理器实例
db_manager = DatabaseManager()


async def get_database_manager() -> DatabaseManager:
    """
    获取数据库管理器实例
    
    Returns:
        DatabaseManager: 数据库管理器实例
    """
    return db_manager


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """
    获取数据库会话（用于FastAPI依赖注入）
    
    Yields:
        AsyncSession: 数据库会话
    """
    if not async_session_factory:
        raise RuntimeError("数据库未初始化")
    
    async with async_session_factory() as session:
        try:
            yield session
            await session.commit()
        except Exception as e:
            await session.rollback()
            logger.error(f"数据库会话错误: {e}", error=str(e))
            raise
        finally:
            await session.close()


@asynccontextmanager
async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
    """
    获取数据库会话上下文管理器
    
    Yields:
        AsyncSession: 数据库会话
    """
    if not async_session_factory:
        raise RuntimeError("数据库未初始化")
    
    async with async_session_factory() as session:
        try:
            yield session
            await session.commit()
        except Exception as e:
            await session.rollback()
            logger.error(f"数据库会话错误: {e}", error=str(e))
            raise


async def init_database() -> None:
    """初始化数据库连接"""
    await db_manager.initialize()


async def close_database() -> None:
    """关闭数据库连接"""
    await db_manager.close()


async def create_tables() -> None:
    """创建数据库表"""
    await db_manager.create_tables()


async def drop_tables() -> None:
    """删除数据库表"""
    await db_manager.drop_tables()


async def database_health_check() -> bool:
    """数据库健康检查"""
    return await db_manager.health_check()


# 数据库事件监听器将在初始化时动态注册


# 导出主要接口
__all__ = [
    "Base",
    "DatabaseManager",
    "db_manager",
    "get_database_manager",
    "get_db",
    "get_db_session",
    "init_database",
    "close_database",
    "create_tables",
    "drop_tables",
    "database_health_check",
]


